import numpy as np
import sympy as sp

from fealpy.pde.linear_elasticity_model import QiModel3d, PolyModel3d, Model2d, HuangModel2d
from linear_elasticity_model2D import GenLinearElasticitymodel2D
from linear_elasticity_model3D import GenLinearElasticitymodel3D 


import matplotlib.pyplot as plt


class PDE():
    def __init__(self,u=None,x=None,Bounadrytype=None,
                    lam=1,mu=0.5):
        '''
        加载pde，2D or 3D
        pde_idx pde问题,边界条件
        '''

        self.u = u
        self.x = x
        if self.x is not None:
            self.gdim = len(x)
        self.lam = lam
        self.mu = mu

    def Displacement_bc(self):
        #位移边界
        if self.gdim == 2:
            pde = GenLinearElasticitymodel2D(self.u,self.x,lam=self.lam,mu=self.mu,
                    Dirichletbd_n='(x0==1)|(x0==0)|(x1==0)|(x1==1)',
                    Dirichletbd_t='(x0==1)|(x0==0)|(x1==0)|(x1==1)')

        elif self.gdim == 3:
            pde = GenLinearElasticitymodel3D(self.u,self.x,lam=self.lam,mu=self.mu,
                        Dirichletbd_n = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))', 
                        Dirichletbd_t0 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))', 
                        Dirichletbd_t1 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))')

        return pde

    def Stress_and_Displacement_bc(self):
        if self.gdim == 2:
            pde = GenLinearElasticitymodel2D(self.u,self.x,lam=self.lam,mu=self.mu,
                Dirichletbd_n='(x0==1)|(x1==0)',Dirichletbd_t='(x0==1)|(x1==0)',
                Neumannbd_nn='(x0==0)|(x1==1)',Neumannbd_nt='(x0==0)|(x1==1)')

        elif self.gdim == 3:
            pde = GenLinearElasticitymodel3D(self.u,self.x,lam=self.lam,mu=self.mu,
                        Dirichletbd_n = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)', 
                        Dirichletbd_t0 = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)', 
                        Dirichletbd_t1 = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)',
                        Neumannbd_nn='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)',
                        Neumannbd_nt0='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)',
                        Neumannbd_nt1='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)')
                        
        return pde




        





